#!/usr/bin/env python
# coding: utf-8

# In[26]:

import torch 
#pip install torch -i https://pypi.tuna.tsinghua.edu.cn/simple
import torch.nn as nn
import pandas as pd
import numpy as np
#pip install numpy -i https://pypi.tuna.tsinghua.edu.cn/simple
import os
import sklearn as sk #pip install scikit-learn -i https://pypi.tuna.tsinghua.edu.cn/simple
from sklearn.model_selection import train_test_split
import torchvision.transforms as transforms  #pip install torchvision==0.2.2.post3
# 将图片转为Tensor格式，才可以训练
import cv2
import tqdm
import time


# In[68]:


# 获得图片路径列表，并且划分训练集和测试集
all_parent_path = []
all_data_list = []
all_label_list = []
# 为了训练需要将字符串格式的标签转换为int形式
label2Digtal = {"circles":0, "squares":1, "triangles":2}
Digtal2Lbel = {0:"circles", 1:"squares", 2:"triangles"}

source_path = r"E:\OneDrive - 西南大学\教学资料\《计算机控制系统》自动化2020级-2022年上课\04-Program\CCSLab\Lab13\shapes"
for p in os.listdir(source_path):
    all_parent_path.append(os.path.join(source_path, p))

#获取所有的图像和对应标签的列表
for img_parent_path in all_parent_path:
    for img in os.listdir(img_parent_path):
        all_data_list.append(os.path.join(img_parent_path, img))
        all_label_list.append(label2Digtal.get(img_parent_path.split("\\")[-1]))


# In[70]:


#所有机器学习和深度学习的必要步骤--划分训练集和测试集
#其中x一般代表图片本身，y代表图片的标签（label）
#参数test_size 代表选用测试集数量占总数据集的比率
x_train, x_test, y_train, y_test = train_test_split(all_data_list, all_label_list, test_size = 0.2)


# In[88]:


#创建数据集类(所有深度学习框架中的必要过程)
class ShapeClsTrainSet(torch.utils.data.Dataset):
    def __init__(self):
        super(ShapeClsTrainSet, self).__init__()
        self.imgs = x_train
        self.labels = y_train
        
    def __getitem__(self, index):
        img = self.imgs[index]
        label = self.labels[index]
        img = cv2.imread(img)
        # 转换图片通道 从 BGR 到 RGB(由于CV2读取图片默认为BGR格式)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        #根据网络模型需要将图像resize到需要的尺寸，此处AlexNet的标准输入尺寸为227*227
        img = cv2.resize(img,(227, 227))
        # 在进行训练和测试之前一定要将图片转换为框架所需要的Tensor格式
        return transforms.ToTensor()(img), label
        
    def __len__(self):
        return len(self.imgs)

class ShapeClsTestSet(torch.utils.data.Dataset):
    def __init__(self):
        super(ShapeClsTestSet, self).__init__()
        self.imgs = x_test
        self.labels = y_test
        
    def __getitem__(self, index):
        img = self.imgs[index]
        label = self.labels[index]
        img = cv2.imread(img)
        # 转换图片通道 从 BGR 到 RGB(由于CV2读取图片默认为BGR格式)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        #根据网络模型需要将图像resize到需要的尺寸，此处AlexNet的标准输入尺寸为227*227
        img = cv2.resize(img,(227, 227))
        # 在进行训练和测试之前一定要将图片转换为框架所需要的Tensor格式
        return transforms.ToTensor()(img), label
        
    def __len__(self):
        return len(self.imgs)


# In[103]:


#实例化创建的数据集
trainSet = ShapeClsTrainSet()
testSet = ShapeClsTestSet()


# In[104]:


#把实例化好的数据集变成DataLoader形式
train_loader = torch.utils.data.DataLoader(trainSet, batch_size=4, shuffle=True)
test_loader = torch.utils.data.DataLoader(testSet, batch_size=4, shuffle=False)


# In[107]:


#定义网络模型(此处使用经典网络AlexNet)
#我们对AlexNet进行FineTune，改变最后分类层输出数量为3使之契合形状分类数据集
class AlexNet(nn.Module):
    def __init__(self,nums_classify):
        super(AlexNet,self).__init__()
        
        self.nums=nums_classify #分类数
        #5层卷积层
        self.conv=nn.Sequential(
            #第一层：输入：224*224*3  输出：54*54*96
            nn.Conv2d(in_channels=3,out_channels=96,kernel_size=11,stride=4,padding=0),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3,stride=2),#最大池化  输入：54*54*96 输出：26*26*96
            #第二层：输入26*26*96
            nn.Conv2d(in_channels=96,out_channels=256,kernel_size=5,stride=1,padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3,stride=2),#最大池化   输入：26*26*256  输出：12*12*256
            #第三层：
            nn.Conv2d(in_channels=256,out_channels=384,kernel_size=3,stride=1,padding=1),
            #第四层：
            nn.Conv2d(in_channels=384,out_channels=256,kernel_size=3,stride=1,padding=1),
            nn.ReLU(),
            #第五层：
            nn.Conv2d(in_channels=256,out_channels=256,kernel_size=3,stride=1,padding=1),
            nn.ReLU(),
            # 池化层  输入：12*12*256    输出：5*5*256
            nn.MaxPool2d(kernel_size=3, stride=2)
        )
        #定义池化层，规定输出size
        self.avg_pool=nn.AdaptiveAvgPool2d((5,5))
        #3层全连接层
        self.fc=nn.Sequential(
            #第6层：
            #输入：5*5*256 -> 4096
            nn.Linear(in_features=5 * 5 * 256, out_features=4096),
            nn.ReLU(),
            nn.Dropout(),
            #第7层
            nn.Linear(in_features=4096, out_features=4096),
            nn.ReLU(),
            nn.Dropout(),
            #第8层 output
            nn.Linear(in_features=4096, out_features=self.nums)
        )

    def forward(self, x):
        x = self.conv(x)
        x = self.avg_pool(x)
        x = x.view(x.size(0), -1)
        #print(x.size())
        x = self.fc(x)
        return x


# In[105]:


#定义训练函数(注意此处为一个Epoch内的执行逻辑编写)
def train_one_epoch(epoch, model, loss_fn, optimizer, train_loader, device=None, scheduler=None):
    model.train() #训练前打开训练模式，测试前打开测试模式，防止flowchart出问题
    
    t = time.time()
    running_loss = 0 #定义一个存放sumloss的变量
    
    #pbar = tqdm(enumerate(train_loader), total=len(train_loader)) # 创建dataloader的进度条对象，方便观察训练过程
    for step, (imgs, img_label) in enumerate(train_loader): # 从loader中取一次(batch_size个)数据即为一个step，每次取出的图片和标签放在一个元组中
        '''
            下面两行注释用于GPU训练的常见写法，CPU训练无需添加(将所需数据送入GPU中)
                    imgs = imgs.to(device).float() #将图片数据放置到gpu上，并且规定为float格式
                    img_label = img_label.to(device).long() #将标签数据放置到gpu上，并且规定为long格式
        '''
        img_preds = model(imgs)
        loss = loss_fn(img_preds, img_label)
        
        optimizer.zero_grad() #pytorch的特点，在进行梯度更新前必须要将梯度置零
        
        loss.backward() #进行反向传播
        optimizer.step() #反向传播后进行梯度更新
        running_loss += loss.item()
    
    print(f"epoch={epoch},loss:{running_loss}")

#定义推理函数(注意此处为一个Epoch内的执行逻辑编写)
def test_fn(model, test_loader, device=None):
    model.eval()
    
    t = time.time()
    img_preds_all=[]
    
    #pbar = tqdm(enumerate(test_loader), total=len(test_loader))
    for step, (imgs, img_labels) in enumerate(test_loader):
        '''
            imgs = imgs.to(device).float() 同样为GPU训练所需要的代码
        '''
        img_preds = model(imgs)
        #img_preds_all += [torch.softmax(image_preds, 1).detach().cpu().numpy()]
        img_preds_all += [torch.argmax(img_preds, 1)]
        
    img_preds_all = np.concatenate(img_preds_all, axis = 0)
    
    return img_preds_all


# In[110]:


#训练缓存未及时清理造成内存问题的时候，可以尝试运行清缓存指令
torch.cuda.empty_cache()


# In[111]:


#主函数
if __name__ =="__main__":
    # 实例化
    model = AlexNet(3)
    # 实例化优化器，并且注意输入需要学习的参数，此处我们需要学习的参数就是FineTune-ALexNet的所有参数
    optimizer = torch.optim.Adam(model.parameters())
    # 定义损失函数-分类常用为交叉熵损失函数
    loss_fn = nn.CrossEntropyLoss()
    # 此处的30为学习的epoch数量,即为训练总轮数
    for epoch in range(30):
        train_one_epoch(epoch, model, loss_fn, optimizer, train_loader)
    #推理过程
    preds = test_fn(model, test_loader)
    #计算推理准确度
    print(f"形状分类模型的准确率为{sum(y_test == preds) / len(y_test)}")




